from typing import Annotated

from langchain_openai import AzureChatOpenAI
from typing_extensions import TypedDict

from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langchain_community.tools.tavily_search import TavilySearchResults

from langchain_core.messages import ToolMessage
from langgraph.prebuilt import ToolNode, tools_condition

from langgraph.checkpoint.memory import MemorySaver


import json
import os
os.environ["TAVILY_API_KEY"] = "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"

# 创建记忆

memory = MemorySaver()

# 工具
tool = TavilySearchResults(max_results=2)
tools = [tool]


params = {"azure_endpoint": "https://azs-dev-us-01.openai.azure.com/",
          "openai_api_key": "73cf593826a54526bebd341aa0ca551e", "model_name": "gpt-3.5-turbo",
          "deployment_name": "gpt-4o",
          "openai_api_version": "2023-07-01-preview", "temperature": 0, "max_tokens": 4096, "top_p": 0,
          "frequency_penalty": 0, "presence_penalty": 0, "streaming": False,
          "request_timeout": 600,
          "max_retries": 0}
class State(TypedDict):
    messages: Annotated[list, add_messages]
graph_builder = StateGraph(State)
# 自定义工具
# class BasicToolNode:
#     """A node that runs the tools requested in the last AIMessage."""
#
#     def __init__(self, tools: list) -> None:
#         self.tools_by_name = {tool.name: tool for tool in tools}
#         print(self.tools_by_name)
#
#     def __call__(self, inputs: dict):
#         if messages := inputs.get("messages", []):
#             message = messages[-1]
#         else:
#             raise ValueError("No message found in input")
#         outputs = []
#         for tool_call in message.tool_calls:
#             tool_result = self.tools_by_name[tool_call["name"]].invoke(
#                 tool_call["args"]
#             )
#             outputs.append(
#                 ToolMessage(
#                     content=json.dumps(tool_result),
#                     name=tool_call["name"],
#                     tool_call_id=tool_call["id"],
#                 )
#             )
#         return {"messages": outputs}

tool_node = ToolNode(tools)



llm = AzureChatOpenAI(
    **params
).bind_tools(tools)




def chatbot(state: State):
    return {"messages": [llm.invoke(state["messages"])]}



graph_builder.add_node("tools", tool_node)
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_conditional_edges('chatbot',tools_condition)
graph_builder.add_edge('tools','chatbot')
graph_builder.set_entry_point("chatbot")
graph_builder.set_finish_point("chatbot")
graph = graph_builder.compile(checkpointer=memory,interrupt_before=["tools"])
# 当前ID 为1的状态

# user_input = "猪八戒是谁?"
# config = {"configurable": {"thread_id": "1"}}
#
# events = graph.stream(
#     {"messages": [("user", user_input)]}, config, stream_mode="values"
# )
# # 事件的执行
# # for event in events:
# #     event["messages"][-1].pretty_print()
#
# # 从中断的地方继续
# # events = graph.stream(None, config, stream_mode="values")
# for event in events:
#     if "messages" in event:
#         event["messages"][-1].pretty_print()



user_input = "周杰伦"
config = {"configurable": {"thread_id": "1"}}  # we'll use thread_id = 2 here
events = graph.stream(
    {"messages": [("user", user_input)]}, config, stream_mode="values"
)
for event in events:
    if "messages" in event:
        event["messages"][-1].pretty_print()

# 更新工具调用
from langchain_core.messages import AIMessage
# # 当前工具的状态
snapshot = graph.get_state(config)
existing_message = snapshot.values["messages"][-1]
print(existing_message.tool_calls[0])
new_tool_call = existing_message.tool_calls[0].copy()
new_tool_call["args"]["query"] = "LangGraph human-in-the-loop workflow"
new_message = AIMessage(
    content=existing_message.content,
    tool_calls=[new_tool_call],
    # Important! The ID is how LangGraph knows to REPLACE the message in the state rather than APPEND this messages
    id=existing_message.id,
)
graph.update_state(config, {"messages": [new_message]})

print("\n\nTool calls")
print(graph.get_state(config).values["messages"][-1].tool_calls)

# 手动更新状态 让他返回要输出的答案
# from langchain_core.messages import AIMessage, ToolMessage
# answer = (
#     "周杰伦是个好父亲,他的歌非常的出色"
# )
# new_messages = [
#     ToolMessage(content=answer, tool_call_id=existing_message.tool_calls[0]["id"]),
#     AIMessage(content=answer),
# ]
# new_messages[-1].pretty_print()
# graph.update_state(
#     config,
#     {"messages": new_messages},
# )
#
# print("\n\nLast 2 messages;")
# print(graph.get_state(config).values["messages"][-2:])
# def stream_graph_updates(user_input: str):
#     for event in graph.stream({"messages": [("user", user_input)]}, config, stream_mode="values"):
#         event["messages"][-1].pretty_print()
#         # for value in event.values():
#         #     print("Assistant:", value["messages"][-1]["content"])
#
#
# while True:
#     try:
#         user_input = input("User: ")
#         if user_input.lower() in ["quit", "exit", "q"]:
#             print("Goodbye!")
#             break
#
#         stream_graph_updates(user_input)
#     except:
#         # fallback if input() is not available
#         user_input = "What do you know about LangGraph?"
#         print("User: " + user_input)
#         stream_graph_updates(user_input)
#         break